import os

import numpy as np
import torch
from math import sqrt

import adjustment_curve
import longitudinal_graph
from config import point_data_root as base_path


class NotACircleError(Exception):
    def __init__(self, arg='not a circle'):
        self.args = arg


def get_x_y_z(s):
    x_y = adjustment_curve.getxy(s,1)[0]
    x, y = x_y[0], x_y[1]
    z = longitudinal_graph.getZ(s,1)
    return x, y, z
# 根据里程得到点的x、y、z坐标


def load_files(path, s):
    dirs = os.listdir(path)
    dirs_num = [float(x) for x in dirs]
    dirs_s = [x for x in dirs_num if (s + 1) >= x >= (s - 1)]
    dirs_char = [str(x) for x in dirs_s]
    return dirs_char
#找出所有距离里程1m内的点云数据文件名，
#dis_char为包含这些文件名的列表


def load_point_data(path, file_names):
    tensors = []
    for file in file_names:
        file_path = os.path.join(path, file)
        tensor = torch.load(file_path).numpy()
        tensors.append(tensor)

    return tensors
#将包含在dis_char中的文件的点云坐标读出来，
#tensors为一个多行三列的tensor，三列分别表示点云的x,y,z坐标

class CalculateOperation:
    def __init__(self, x0, y0, z0, xleft, yleft, zleft, xright, yright, zright):
        self.x0 = x0
        self.y0 = y0
        self.z0 = z0
        self.A = xright - xleft
        self.B = yright - yleft
        self.C = zright - zleft
        self.D = -x0 * self.A - y0 * self.B - z0 * self.C
        # A,B,C,D为法平面方程对应参数

    # 求x轴方向向量，注意考虑平行于三维坐标系y轴的情况
    def get_x_axis(self):
        if self.B != 0:
            y_pos = (-self.C * self.z0 - self.D - self.A * (self.x0 + 1)) / self.B - self.y0
            self.x_axis = np.array([1, y_pos, 0])
        else:
            self.x_axis = np.array([0., 1., 0.])

    # 将平面法向量与二维坐标系下x轴的方向向量叉乘，得到y轴方向向量
    def get_y_axis(self):
        self.y_axis = np.cross(self.x_axis, np.array([self.A, self.B, self.C]))
        if self.y_axis[2] < 0:
            self.y_axis = -self.y_axis
            # 此处默认求得方向向量的z大于0，
            # 即若叉乘后 y = [1,2,-1]，则令 y = [-1,-2,1]

    # 求点到已知平面距离
    def get_distance_point_surface(self, x, y, z):
        return np.abs(self.A * x + self.B * y + self.C * z + self.D) / (
            np.sqrt(self.A ** 2 + self.B ** 2 + self.C ** 2))

    def test(self, points):
        # 先传入点云坐标points
        self.points = points
        # 求点云到法平面距离,并取出距离小于0.1m的点
        points_ditance_surface = self.get_distance_point_surface(points[:, 0], points[:, 1], points[:, 2])
        self.points = self.points[points_ditance_surface < 0.05]
        #print("num:",self.points.shape[0])
        self.l = self.points.shape[0]
        #if self.l < 100:
            #raise NotACircleError("only few points")

        #         fig = plt.figure(figsize=(15, 4))
        #         ax = plt.gca(projection='3d')
        #         ax.scatter(self.points[:,0], self.points[:,1], self.points[:,2])

        # 求二维坐标系x轴和y轴的方向向量
        self.get_x_axis()
        self.get_y_axis()
        # 将这两个方向向量转化为单位向量
        self.x_axis /= np.linalg.norm(self.x_axis)
        self.y_axis /= np.linalg.norm(self.y_axis)

        # 返回横纵坐标
        i = self.points - np.array([self.x0, self.y0, self.z0])
        return np.sum(self.x_axis * i, axis=1), np.sum(self.y_axis * i, axis=1)


def init_object(s):
    x0, y0, z0 = get_x_y_z(s)
    xleft, yleft, zleft = get_x_y_z(s - 0.01)
    xright, yright, zright = get_x_y_z(s + 0.01)
#用原里程加减1mm来构造一条直线以求取法平面

    return CalculateOperation(x0, y0, z0, xleft, yleft, zleft, xright, yright, zright)
#构造对象

def get_points(s):
    dis_char = load_files(base_path, s)#得到包含所有文件名的列表
    points = load_point_data(base_path, dis_char)
    points = np.concatenate(points)
    #读取出范围内的所有点云的三维坐标
    return points


def least_square(x_cor, y_cor):
    #最小二乘法拟合圆
    x_cube_sum = np.sum(x_cor ** 3)
    y_cube_sum = np.sum(y_cor ** 3)
    x_square_sum = np.sum(x_cor ** 2)
    y_square_sum = np.sum(y_cor ** 2)
    x_sum = np.sum(x_cor)
    y_sum = np.sum(y_cor)
    x_times_y_sum = np.sum(x_cor * y_cor)

    N = x_cor.size

    C = N * x_square_sum - x_sum * x_sum
    D = N * x_times_y_sum - x_sum * y_sum
    E = N * x_cube_sum + N * np.sum(x_cor * (y_cor ** 2)) - np.sum(x_cor ** 2 + y_cor ** 2) * x_sum
    G = N * y_square_sum - y_sum * y_sum
    H = N * np.sum(x_cor ** 2 * y_cor) + N * y_cube_sum - np.sum(x_cor ** 2 + y_cor ** 2) * y_sum
    a = (H * D - E * G) / (C * G - D ** 2)
    b = (H * C - E * D) / (D ** 2 - G * C)
    c = -(np.sum(x_cor ** 2 + y_cor ** 2) + a * x_sum + b * y_sum) / N
    A = a / (-2)
    B = b / (-2)
    R = np.sqrt(a ** 2 + b ** 2 - 4 * c) / 2

    return -A, B, R


def main_function(s):
    #创建对象
    cal = init_object(s)
    #求点云投影后转化得到的二维坐标
    x_cor, y_cor = cal.test(get_points(s))
    #     plt.scatter(x_cor, y_cor)
    #用最小二乘法拟合圆的方程，返回圆心坐标和半径
    A, B, R = least_square(x_cor, y_cor)
    if R<2.6 or R>2.8:
        raise NotACircleError
    return x_cor, y_cor, A, B, R


def calculate_coordinate_by_x(s, x):
    x_cor, y_cor, A, B, R = main_function(s)
    A, B, R = A*1000, B*1000, R*1000

    # flag = adjustment_curve.getxy(s)[-1]
    # A *= flag
    A *= -1
    B = longitudinal_graph.getZ(s,1) * 1000 + B
    x *= 1000
    # print('(x - {})^2 + (y - {})^2 = {}'.format(A, B, R * R))
    return [sqrt(R * R - pow((x - A), 2)) + B, -sqrt(R * R - pow((x - A), 2)) + B]


def calculate_coordinate_by_y(s, y, add_Z=True):
    x_cor, y_cor, A, B, R = main_function(s)
    A, B, R = A*1000, B*1000, R*1000
    # print("所拟合的圆圆心坐标为(", A, ",", B, ")")
    # print("半径为：", R)

    # flag = adjustment_curve.getxy(s)[-1]
    # A *= (-flag if flag > 0 else flag)
    A *= -1
    y = (y - longitudinal_graph.getZ(s,1) if add_Z else 0) * 1000
    # print('(x - {})^2 + (y - {})^2 = {}'.format(A, B, R * R))
    return [sqrt(R * R - pow((y - B), 2)) + A, -sqrt(R * R - pow((y - B), 2)) + A]


if __name__ == '__main__':
    # print(main_function(-1500))
    s = -1583.506  # 里程
    y = -8.11  # 点的y坐标
    x1, x2 = calculate_coordinate_by_y(s, y)
    y1, y2 = calculate_coordinate_by_x(s, 0)  # 根据x计算y
    print(x1, x2, y1, y2)
